import tensorflow as tf
from polyaxon.tracking import Run

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

experiment = Run(project='mnist', artifacts_path='/tmp/mnist/')
experiment.create(tags=['examples', 'tensorflow'])


def create_model():
    return tf.keras.models.Sequential([
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(512, activation='relu'),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10, activation='softmax')
    ])


model = create_model()
model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])


model.fit(x=x_train,
          y=y_train,
          epochs=5,
          validation_data=(x_test, y_test))
